import sys
import numpy as np
sys.path.append(sys.argv[3])
from signal_udfs import lowpass
import matplotlib.pyplot as plt
import pickle

input_file = sys.argv[1]
output_handle = sys.argv[2]

#Import photometry data from Doric .csv file
time = []
gcamp_channel = []
control_channel = []
ttl_channel = []

for line_num, line in enumerate(open(input_file)):
	print('Importing data... 0 % Complete', end = '\r')
	line = line.strip().split(',')
	if line_num >= 1 and line[1] != '' and float(line[0]) >= 0.2:
		time.append(float(line[0]))
		gcamp_channel.append(float(line[1]))
		control_channel.append(float(line[2]))

print('Importing data... 100 % Complete')

print('Calculating DF/F0... 0 % Complete', end ='\r')
			
#Determine sampling rate
fs_counter = 0
fs_found = 0
start_ts = time[1]

for x in time:
	fs_counter += 1
	if x - start_ts >= 1:
		break

#3 Hz Butterwort lowpass filter channels
gcamp_channel_filt = lowpass(gcamp_channel, 3, fs_counter, 2)
control_channel_filt = lowpass(control_channel, 3, fs_counter, 2)

#Fit the control channel to the gcamp channel
line_fit = np.polyfit(control_channel_filt, gcamp_channel_filt, 1)

a = line_fit[0]
b = line_fit[1]

control_fit = a * control_channel_filt + b

#Calculate DF/F0
norm_dat = (gcamp_channel_filt - control_fit) / control_fit
norm_dat = norm_dat * 100

#Output block
df_f_output = open(output_handle[:-4] + '_df-f0.csv','w')
title_line = ['Time (s)', '465nm', '405nm', '405nm Fitted', 'DF/F0']

df_f_output.write(','.join(title_line)+'\n')

for item_num, item in enumerate(time):
	current_line = []
	current_line.append(str(item))
	current_line.append(str(gcamp_channel_filt[item_num]))
	current_line.append(str(control_channel_filt[item_num]))
	current_line.append(str(control_fit[item_num]))
	current_line.append(str(norm_dat[item_num]))
	df_f_output.write(','.join(current_line) + '\n')
	
#Plot DF/F0 and save figure
plt.plot(time,norm_dat)
plt.xlabel('Time (s)')
plt.ylabel('DF/F0')
plt.savefig(output_handle[:-4] + '_df-f0.png')
#plt.show()
plt.close()

#Plot and save 465nm and fitted 405nm channels
plt.plot(time, control_fit, label = '405nm Fit')
plt.plot(time, gcamp_channel_filt, label = '465nm')
plt.xlabel('Time (s)')
plt.ylabel('Volts')
plt.legend(bbox_to_anchor=(1.02,1), loc="upper left")
plt.savefig(output_handle[:-4] + '_channels.png', bbox_inches = 'tight')
#plt.show()
plt.close()

#Pickle base data
pickle.dump(time,open('time.pkl','wb'))
pickle.dump(norm_dat,open('df-f.pkl','wb'))
pickle.dump(ttl_channel,open('ttl.pkl','wb'))

print('Calculating DF/F0... 100 % Complete')
